from typing import Dict, Any, Optional, List, Tuple, Callable
from datetime import datetime, timedelta
import asyncio
import json
from dataclasses import dataclass, asdict
from enum import Enum
import asyncpg
from sqlalchemy import Column, String, Text, DateTime, Float, Boolean, Integer, func, select
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.dialects.postgresql import UUID, JSONB
import uuid
import time

from utils.logger import Logger, get_logger
from utils.exceptions import DatabaseError
from middleware.data_consistency import data_consistency_manager, transaction

logger = get_logger("database_manager")

Base = declarative_base()


class TaskStatus(Enum):
    """任务状态"""
    PENDING = "pending"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class AgentType(Enum):
    """智能体类型"""
    ORCHESTRATOR = "orchestrator"
    TEXT_RETRIEVAL = "text_retrieval"
    DOC_SEGMENT = "doc_segment"
    MULTI_MODAL = "multi_modal"
    VALIDATION = "validation"
    INTEGRATION = "integration"
    FALLBACK = "fallback"
    MONITOR = "monitor"


# 数据库模型
class TaskRecord(Base):
    """任务记录表"""
    __tablename__ = "task_records"
    
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    task_id = Column(String(255), unique=True, nullable=False, index=True)
    agent_type = Column(String(50), nullable=False, index=True)
    task_type = Column(String(50), nullable=False)
    status = Column(String(20), nullable=False, index=True)
    query = Column(Text)
    input_data = Column(JSONB)
    output_data = Column(JSONB)
    error_message = Column(Text)
    execution_time = Column(Float)
    retry_count = Column(Integer, default=0)
    created_at = Column(DateTime, default=datetime.utcnow, index=True)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    completed_at = Column(DateTime)


class AgentMetrics(Base):
    """智能体指标表"""
    __tablename__ = "agent_metrics"
    
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    agent_type = Column(String(50), nullable=False, index=True)
    metric_name = Column(String(100), nullable=False)
    metric_value = Column(Float, nullable=False)
    labels = Column(JSONB)
    timestamp = Column(DateTime, default=datetime.utcnow, index=True)


class DocumentMetadata(Base):
    """文档元数据表"""
    __tablename__ = "document_metadata"
    
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    doc_id = Column(String(255), unique=True, nullable=False, index=True)
    filename = Column(String(500), nullable=False)
    file_path = Column(String(1000))
    file_type = Column(String(50))
    file_size = Column(Integer)
    metadata = Column(JSONB)
    content_hash = Column(String(64), index=True)
    processing_status = Column(String(20), default="pending")
    created_at = Column(DateTime, default=datetime.utcnow, index=True)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)


class DocumentSegment(Base):
    """文档分段表"""
    __tablename__ = "document_segments"
    
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    doc_id = Column(String(255), nullable=False, index=True)
    segment_id = Column(String(255), nullable=False, index=True)
    segment_type = Column(String(50))  # chapter, section, paragraph
    title = Column(String(500))
    content = Column(Text)
    page_start = Column(Integer)
    page_end = Column(Integer)
    position_start = Column(Integer)
    position_end = Column(Integer)
    metadata = Column(JSONB)
    vector_embedding = Column(JSONB)  # 存储向量嵌入
    created_at = Column(DateTime, default=datetime.utcnow, index=True)


class SystemConfig(Base):
    """系统配置表"""
    __tablename__ = "system_config"
    
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    config_key = Column(String(100), unique=True, nullable=False, index=True)
    config_value = Column(JSONB)
    description = Column(Text)
    is_active = Column(Boolean, default=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)


@dataclass
class DatabaseConfig:
    """数据库配置"""
    host: str = "192.168.64.100"
    port: int = 5432
    database: str = "maess_flow"
    username: str = "maess_user"
    password: str = "maess_password"
    pool_size: int = 20
    max_overflow: int = 30
    pool_timeout: int = 30
    echo_sql: bool = False


class DatabaseManager:
    """PostgreSQL数据库管理器"""

    def __init__(self, config: Optional[DatabaseConfig] = None):
        self.config = config or DatabaseConfig()
        self.engine = None
        self.async_pool = None
        self.session_factory = None
        self._initialize_database()

    def _initialize_database(self):
        """
        初始化数据库连接
        """
        try:
            # 创建异步引擎
            async_url = f"postgresql+asyncpg://{self.config.username}:{self.config.password}@{self.config.host}:{self.config.port}/{self.config.database}"
            
            self.engine = create_async_engine(
                async_url,
                pool_size=self.config.pool_size,
                max_overflow=self.config.max_overflow,
                pool_timeout=self.config.pool_timeout,
                echo=self.config.echo_sql
            )
            
            # 创建异步会话工厂
            self.async_session_factory = async_sessionmaker(
                bind=self.engine,
                class_=AsyncSession,
                autocommit=False,
                autoflush=False,
                expire_on_commit=False
            )
            
            logger.info("数据库初始化完成")
            
        except Exception as e:
            logger.error(f"数据库初始化失败: {str(e)}")
            raise DatabaseError(f"数据库初始化失败: {str(e)}")
    
    async def create_tables(self):
        """
        异步创建数据库表
        """
        try:
            async with self.engine.begin() as conn:
                await conn.run_sync(Base.metadata.create_all)
            logger.info("数据库表创建成功")
        except Exception as e:
            logger.error(f"数据库表创建失败: {str(e)}")
            raise DatabaseError(f"数据库表创建失败: {str(e)}")

    async def create_task_record(self, task_id: str, agent_type: str, 
                                task_type: str, query: str, 
                                input_data: Dict[str, Any]) -> str:
        """
        创建任务记录
        """
        try:
            async with self.async_session_factory() as session:
                async with session.begin():
                    task_record = TaskRecord(
                        task_id=task_id,
                        agent_type=agent_type,
                        task_type=task_type,
                        status=TaskStatus.PENDING.value,
                        query=query,
                        input_data=input_data
                    )
                    
                    session.add(task_record)
                    await session.flush()
                    
                    record_id = str(task_record.id)
                    logger.info(f"任务记录创建完成: task_id={task_id}, record_id={record_id}")
                    
                    return record_id
            
        except Exception as e:
            logger.error(f"创建任务记录失败: {str(e)}")
            raise DatabaseError(f"创建任务记录失败: {str(e)}")

    async def update_task_status(self, task_id: str, status: TaskStatus,
                               output_data: Optional[Dict[str, Any]] = None,
                               error_message: Optional[str] = None,
                               execution_time: Optional[float] = None) -> bool:
        """
        更新任务状态
        """
        try:
            async with self.async_session_factory() as session:
                async with session.begin():
                    stmt = select(TaskRecord).where(TaskRecord.task_id == task_id)
                    result = await session.execute(stmt)
                    task_record = result.scalar_one_or_none()
                    
                    if not task_record:
                        raise DatabaseError(f"任务记录不存在: task_id={task_id}")
                    
                    task_record.status = status.value
                    
                    if output_data:
                        task_record.output_data = output_data
                    
                    if error_message:
                        task_record.error_message = error_message
                    
                    if execution_time:
                        task_record.execution_time = execution_time
                    
                    if status in [TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED]:
                        task_record.completed_at = datetime.utcnow()
                
            logger.info(f"任务状态更新完成: task_id={task_id}, status={status.value}")
            return True
            
        except Exception as e:
            logger.error(f"更新任务状态失败: {str(e)}")
            raise DatabaseError(f"更新任务状态失败: {str(e)}")

    async def get_task_record(self, task_id: str) -> Optional[Dict[str, Any]]:
        """
        获取任务记录
        """
        try:
            async with self.async_session_factory() as session:
                stmt = select(TaskRecord).where(TaskRecord.task_id == task_id)
                result = await session.execute(stmt)
                task_record = result.scalar_one_or_none()
                
                if not task_record:
                    return None
                
                return {
                    "id": str(task_record.id),
                    "task_id": task_record.task_id,
                    "agent_type": task_record.agent_type,
                    "task_type": task_record.task_type,
                    "status": task_record.status,
                    "query": task_record.query,
                    "input_data": task_record.input_data,
                    "output_data": task_record.output_data,
                    "error_message": task_record.error_message,
                    "execution_time": task_record.execution_time,
                    "retry_count": task_record.retry_count,
                    "created_at": task_record.created_at.isoformat() if task_record.created_at else None,
                    "updated_at": task_record.updated_at.isoformat() if task_record.updated_at else None,
                    "completed_at": task_record.completed_at.isoformat() if task_record.completed_at else None
                }
            
        except Exception as e:
            logger.error(f"获取任务记录失败: {str(e)}")
            raise DatabaseError(f"获取任务记录失败: {str(e)}")

    async def get_task_records(self, agent_type: Optional[str] = None,
                             status: Optional[TaskStatus] = None,
                             limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]:
        """
        获取任务记录列表
        """
        try:
            async with self.async_session_factory() as session:
                stmt = select(TaskRecord)
                
                if agent_type:
                    stmt = stmt.where(TaskRecord.agent_type == agent_type)
                
                if status:
                    stmt = stmt.where(TaskRecord.status == status.value)
                
                stmt = stmt.order_by(TaskRecord.created_at.desc()).limit(limit).offset(offset)
                result = await session.execute(stmt)
                records = result.scalars().all()
                
                result_list = []
                for record in records:
                    result_list.append({
                        "id": str(record.id),
                        "task_id": record.task_id,
                        "agent_type": record.agent_type,
                        "task_type": record.task_type,
                        "status": record.status,
                        "query": record.query,
                        "execution_time": record.execution_time,
                        "created_at": record.created_at.isoformat() if record.created_at else None,
                        "completed_at": record.completed_at.isoformat() if record.completed_at else None
                    })
                
                return result_list
            
        except Exception as e:
            logger.error(f"获取任务记录列表失败: {str(e)}")
            raise DatabaseError(f"获取任务记录列表失败: {str(e)}")

    async def save_agent_metrics(self, agent_type: str, metric_name: str,
                               metric_value: float, labels: Optional[Dict[str, Any]] = None) -> str:  
        """保存智能体指标"""
        try:
            async with self.async_session_factory() as session:
                async with session.begin():
                    metrics_record = AgentMetrics(
                        agent_type=agent_type,
                        metric_name=metric_name,
                        metric_value=metric_value,
                        labels=labels or {}
                    )
                    
                    session.add(metrics_record)
                    await session.flush()
                    
                    record_id = str(metrics_record.id)
                    logger.debug(f"智能体指标保存完成: agent_type={agent_type}, metric_name={metric_name}")
                    
                    return record_id
                    
        except Exception as e:
            logger.error(f"保存智能体指标失败: {str(e)}")
            raise DatabaseError(f"保存智能体指标失败: {str(e)}")

    async def get_agent_metrics(self, agent_type: str, metric_name: str,
                              time_range_hours: int = 24) -> List[Dict[str, Any]]:
        """获取智能体指标"""
        try:
            async with self.async_session_factory() as session:
                from_time = datetime.utcnow() - timedelta(hours=time_range_hours)
                
                stmt = select(AgentMetrics).where(
                    AgentMetrics.agent_type == agent_type,
                    AgentMetrics.metric_name == metric_name,
                    AgentMetrics.timestamp >= from_time
                ).order_by(AgentMetrics.timestamp.desc())
                
                result = await session.execute(stmt)
                records = result.scalars().all()
                
                metrics_list = []
                for record in records:
                    metrics_list.append({
                        "id": str(record.id),
                        "metric_value": record.metric_value,
                        "labels": record.labels,
                        "timestamp": record.timestamp.isoformat() if record.timestamp else None
                    })
                
                return metrics_list
                
        except Exception as e:
            logger.error(f"获取智能体指标失败: {str(e)}")
            raise DatabaseError(f"获取智能体指标失败: {str(e)}")

    async def save_document_metadata(self, doc_id: str, filename: str,
                                   file_path: Optional[str] = None,
                                   file_type: Optional[str] = None,
                                   file_size: Optional[int] = None,
                                   metadata: Optional[Dict[str, Any]] = None,
                                   content_hash: Optional[str] = None,
                                   transaction_id: Optional[str] = None) -> str:
        """
        保存文档元数据，支持事务一致性
        
        :param doc_id: 文档ID
        :param filename: 文件名
        :param file_path: 文件路径
        :param file_type: 文件类型
        :param file_size: 文件大小
        :param metadata: 文档元数据
        :param content_hash: 内容哈希值
        :param transaction_id: 事务ID，如果提供，则使用数据一致性管理器
        :return: 记录ID
        """
        try:
            async with self.async_session_factory() as session:
                async with session.begin():
                    doc_record = DocumentMetadata(
                        doc_id=doc_id,
                        filename=filename,
                        file_path=file_path,
                        file_type=file_type,
                        file_size=file_size,
                        metadata=metadata or {},
                        content_hash=content_hash
                    )
                    
                    session.add(doc_record)
                    await session.flush()  # 立即执行以获取ID
                    
                    # 记录操作信息用于事务恢复
                    if transaction_id:
                        data_consistency_manager.add_operation(transaction_id, {
                            "type": "save_document_metadata",
                            "doc_id": doc_id,
                            "metadata_id": str(doc_record.id)
                        })
                    
                    record_id = str(doc_record.id)
                    logger.info(f"文档元数据保存完成: doc_id={doc_id}, filename={filename}")
                    
                    return record_id
                    
        except Exception as e:
            logger.error(f"保存文档元数据失败: {str(e)}")
            
            # 记录失败信息
            if transaction_id:
                data_consistency_manager.add_operation(transaction_id, {
                    "type": "save_document_metadata_failed",
                    "doc_id": doc_id,
                    "error": str(e)
                })
                
            raise DatabaseError(f"保存文档元数据失败: {str(e)}")

    async def get_document_metadata(self, doc_id: str) -> Optional[Dict[str, Any]]:
        """获取文档元数据"""
        try:
            async with self.async_session_factory() as session:
                stmt = select(DocumentMetadata).where(DocumentMetadata.doc_id == doc_id)
                result = await session.execute(stmt)
                doc_record = result.scalar_one_or_none()
                
                if not doc_record:
                    return None
                
                return {
                    "id": str(doc_record.id),
                    "doc_id": doc_record.doc_id,
                    "filename": doc_record.filename,
                    "file_path": doc_record.file_path,
                    "file_type": doc_record.file_type,
                    "file_size": doc_record.file_size,
                    "metadata": doc_record.metadata,
                    "content_hash": doc_record.content_hash,
                    "processing_status": doc_record.processing_status,
                    "created_at": doc_record.created_at.isoformat() if doc_record.created_at else None,
                    "updated_at": doc_record.updated_at.isoformat() if doc_record.updated_at else None
                }
                
        except Exception as e:
            logger.error(f"获取文档元数据失败: {str(e)}")
            raise DatabaseError(f"获取文档元数据失败: {str(e)}")

    async def save_document_segment(self, doc_id: str, segment_id: str,
                                  segment_type: str, title: Optional[str] = None,
                                  content: Optional[str] = None,
                                  page_start: Optional[int] = None,
                                  page_end: Optional[int] = None,
                                  position_start: Optional[int] = None,
                                  position_end: Optional[int] = None,
                                  metadata: Optional[Dict[str, Any]] = None,
                                  vector_embedding: Optional[List[float]] = None,
                                  transaction_id: Optional[str] = None) -> str:
        """
        保存文档分段，支持事务一致性
        
        :param doc_id: 文档ID
        :param segment_id: 分段ID
        :param segment_type: 分段类型
        :param title: 分段标题
        :param content: 分段内容
        :param page_start: 起始页码
        :param page_end: 结束页码
        :param position_start: 起始位置
        :param position_end: 结束位置
        :param metadata: 分段元数据
        :param vector_embedding: 向量嵌入
        :param transaction_id: 事务ID，如果提供，则使用数据一致性管理器
        :return: 记录ID
        """
        try:
            async with self.async_session_factory() as session:
                async with session.begin():
                    segment_record = DocumentSegment(
                        doc_id=doc_id,
                        segment_id=segment_id,
                        segment_type=segment_type,
                        title=title,
                        content=content,
                        page_start=page_start,
                        page_end=page_end,
                        position_start=position_start,
                        position_end=position_end,
                        metadata=metadata or {},
                        vector_embedding=vector_embedding
                    )
                    
                    session.add(segment_record)
                    await session.flush()  # 立即执行以获取ID
                    
                    # 记录操作信息用于事务恢复
                    if transaction_id:
                        data_consistency_manager.add_operation(transaction_id, {
                            "type": "save_document_segment",
                            "doc_id": doc_id,
                            "segment_id": segment_id,
                            "segment_db_id": str(segment_record.id)
                        })
                    
                    record_id = str(segment_record.id)
                    logger.debug(f"文档分段保存完成: doc_id={doc_id}, segment_id={segment_id}")
                    
                    return record_id
                    
        except Exception as e:
            logger.error(f"保存文档分段失败: {str(e)}")
            
            # 记录失败信息
            if transaction_id:
                data_consistency_manager.add_operation(transaction_id, {
                    "type": "save_document_segment_failed",
                    "doc_id": doc_id,
                    "segment_id": segment_id,
                    "error": str(e)
                })
                
            raise DatabaseError(f"保存文档分段失败: {str(e)}")

    async def get_document_segments(self, doc_id: str,
                                  segment_type: Optional[str] = None) -> List[Dict[str, Any]]:
        """获取文档分段"""
        try:
            async with self.async_session_factory() as session:
                stmt = select(DocumentSegment).where(DocumentSegment.doc_id == doc_id)
                
                if segment_type:
                    stmt = stmt.where(DocumentSegment.segment_type == segment_type)
                
                stmt = stmt.order_by(DocumentSegment.position_start.asc())
                result = await session.execute(stmt)
                segments = result.scalars().all()
                
                segments_list = []
                for segment in segments:
                    segments_list.append({
                        "id": str(segment.id),
                        "segment_id": segment.segment_id,
                        "segment_type": segment.segment_type,
                        "title": segment.title,
                        "content": segment.content,
                        "page_start": segment.page_start,
                        "page_end": segment.page_end,
                        "position_start": segment.position_start,
                        "position_end": segment.position_end,
                        "metadata": segment.metadata,
                        "vector_embedding": segment.vector_embedding,
                        "created_at": segment.created_at.isoformat() if segment.created_at else None
                    })
                
                return segments_list
                
        except Exception as e:
            logger.error(f"获取文档分段失败: {str(e)}")
            raise DatabaseError(f"获取文档分段失败: {str(e)}")

    async def get_system_config(self, config_key: str) -> Optional[Dict[str, Any]]:
        """获取系统配置"""
        try:
            async with self.async_session_factory() as session:
                stmt = select(SystemConfig).where(SystemConfig.config_key == config_key)
                result = await session.execute(stmt)
                config = result.scalar_one_or_none()
                
                if not config:
                    return None
                
                return {
                    "id": str(config.id),
                    "config_key": config.config_key,
                    "config_value": config.config_value,
                    "description": config.description,
                    "is_active": config.is_active,
                    "created_at": config.created_at.isoformat() if config.created_at else None,
                    "updated_at": config.updated_at.isoformat() if config.updated_at else None
                }
                
        except Exception as e:
            logger.error(f"获取系统配置失败: {str(e)}")
            raise DatabaseError(f"获取系统配置失败: {str(e)}")

    async def set_system_config(self, config_key: str, config_value: Dict[str, Any],
                              description: Optional[str] = None) -> str:
        """设置系统配置"""
        try:
            async with self.async_session_factory() as session:
                async with session.begin():
                    stmt = select(SystemConfig).where(SystemConfig.config_key == config_key)
                    result = await session.execute(stmt)
                    config_record = result.scalar_one_or_none()
                    
                    if config_record:
                        config_record.config_value = config_value
                        config_record.description = description
                        config_record.is_active = True
                    else:
                        config_record = SystemConfig(
                            config_key=config_key,
                            config_value=config_value,
                            description=description,
                            is_active=True
                        )
                        session.add(config_record)
                    
                    await session.flush()
                    
                    record_id = str(config_record.id)
                    logger.info(f"系统配置设置完成: config_key={config_key}")
                    
                    return record_id
                    
        except Exception as e:
            logger.error(f"设置系统配置失败: {str(e)}")
            raise DatabaseError(f"设置系统配置失败: {str(e)}")


    async def execute_query(self, query: str, params: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
        """执行自定义SQL查询"""
        try:
            async with self.async_session_factory() as session:
                result = await session.execute(query, params or {})
                
                # 获取列名
                columns = result.keys()
                
                # 构建结果
                rows = []
                for row in result.fetchall():
                    row_dict = {}
                    for i, column in enumerate(columns):
                        value = row[i]
                        if isinstance(value, datetime):
                            value = value.isoformat()
                        elif isinstance(value, uuid.UUID):
                            value = str(value)
                        row_dict[column] = value
                    rows.append(row_dict)
                
                return rows
                
        except Exception as e:
            logger.error(f"执行查询失败: {str(e)}")
            raise DatabaseError(f"执行查询失败: {str(e)}")

    async def health_check(self) -> Dict[str, Any]:
        """健康检查"""
        try:
            start_time = time.time()
            
            # 执行简单查询
            async with self.async_session_factory() as session:
                result = await session.execute("SELECT 1")
                
            # 统计连接池信息
            pool_size = self.engine.pool.size()
            checked_out = self.engine.pool.checkedout()
            
            execution_time = time.time() - start_time
            
            return {
                "status": "healthy",
                "database": "connected",
                "execution_time_ms": round(execution_time * 1000, 2),
                "pool_size": pool_size,
                "checked_out_connections": checked_out,
                "available_connections": pool_size - checked_out,
                "timestamp": datetime.now().isoformat()
            }
            
        except Exception as e:
            logger.error(f"健康检查失败: {str(e)}")
            return {
                "status": "unhealthy",
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            }

    async def save_document_segments(self, segments: List[Dict[str, Any]], 
                                   transaction_id: Optional[str] = None) -> bool:
        """
        批量保存文档分段，支持事务一致性
        
        :param segments: 分段信息列表，每个元素包含分段所需的字段
        :param transaction_id: 事务ID，如果提供，则使用数据一致性管理器
        :return: 是否保存成功
        """
        try:
            async with self.async_session_factory() as session:
                async with session.begin():
                    segment_records = []
                    segment_info = []
                    
                    for segment_data in segments:
                        segment = DocumentSegment(
                            doc_id=segment_data['doc_id'],
                            segment_id=segment_data['segment_id'],
                            segment_type=segment_data['segment_type'],
                            title=segment_data.get('title'),
                            content=segment_data.get('content'),
                            page_start=segment_data.get('page_start'),
                            page_end=segment_data.get('page_end'),
                            position_start=segment_data.get('position_start'),
                            position_end=segment_data.get('position_end'),
                            metadata=segment_data.get('metadata', {}),
                            vector_embedding=segment_data.get('vector_embedding')
                        )
                        segment_records.append(segment)
                        segment_info.append({
                            'doc_id': segment_data['doc_id'],
                            'segment_id': segment_data['segment_id']
                        })
                    
                    session.add_all(segment_records)
                    await session.flush()  # 确保数据被写入数据库
                    
                    # 记录操作信息用于事务恢复
                    if transaction_id:
                        data_consistency_manager.add_operation(transaction_id, {
                            "type": "save_document_segments",
                            "segments": segment_info,
                            "count": len(segments)
                        })
                    
                    logger.info(f"批量保存文档分段成功，数量: {len(segments)}")
                    return True
                    
        except Exception as e:
            logger.error(f"批量保存文档分段失败: {str(e)}")
            
            # 记录失败信息
            if transaction_id:
                data_consistency_manager.add_operation(transaction_id, {
                    "type": "save_document_segments_failed",
                    "count": len(segments),
                    "error": str(e)
                })
                
            raise DatabaseError(f"批量保存文档分段失败: {str(e)}")
    
    async def perform_consistent_document_operation(self, document_id: str, 
                                           operation_type: str, 
                                           operation_func: Callable, 
                                           **kwargs) -> Dict[str, Any]:
        """
        执行一致的文档操作，确保数据库和向量存储的数据一致性
        
        :param document_id: 文档ID
        :param operation_type: 操作类型
        :param operation_func: 实际执行的操作函数
        :param kwargs: 操作函数的参数
        :return: 操作结果
        """
        # 生成事务ID
        transaction_id = f"doc_op_{operation_type}_{document_id}_{int(time.time())}"
        
        try:
            # 使用事务上下文管理器
            with transaction(transaction_id) as tx_manager:
                # 在事务中执行操作
                result = await operation_func(transaction_id=transaction_id, **kwargs)
                
                # 记录事务成功
                logger.info(f"一致的文档操作成功: {operation_type}, 文档ID: {document_id}")
                
                return {
                    "success": True,
                    "result": result,
                    "transaction_id": transaction_id
                }
                
        except Exception as e:
            logger.error(f"一致的文档操作失败: {operation_type}, 文档ID: {document_id}, 错误: {str(e)}")
            return {
                "success": False,
                "error": str(e),
                "transaction_id": transaction_id
            }
    
    async def verify_data_integrity(self, document_id: Optional[str] = None) -> Dict[str, Any]:
        """
        验证数据完整性
        
        :param document_id: 可选的文档ID，如果为None则验证所有文档
        :return: 完整性验证结果
        """
        try:
            async with self.async_session_factory() as session:
                # 构建基础查询
                metadata_stmt = select(DocumentMetadata)
                if document_id:
                    metadata_stmt = metadata_stmt.where(DocumentMetadata.doc_id == document_id)
                
                segment_stmt = select(DocumentSegment)
                if document_id:
                    segment_stmt = segment_stmt.where(DocumentSegment.doc_id == document_id)
                
                # 统计文档元数据和分段数量
                metadata_result = await session.execute(metadata_stmt)
                metadata_records = metadata_result.scalars().all()
                metadata_count = len(metadata_records)
                
                segment_result = await session.execute(segment_stmt)
                segment_records = segment_result.scalars().all()
                segment_count = len(segment_records)
                
                # 检查是否有孤立的分段
                if document_id:
                    # 检查特定文档的分段
                    orphaned_segments = 0
                else:
                    # 检查所有文档的孤立分段
                    # 使用子查询查找所有有效的doc_id
                    valid_doc_ids = select(DocumentMetadata.doc_id).subquery()
                    orphaned_segments_stmt = select(DocumentSegment).where(
                        ~DocumentSegment.doc_id.in_(valid_doc_ids)
                    )
                    orphaned_segments_result = await session.execute(orphaned_segments_stmt)
                    orphaned_segments = len(orphaned_segments_result.scalars().all())
                
                # 检查文档元数据是否都有对应的分段
                if document_id:
                    # 检查特定文档是否有分段
                    doc_has_segments = segment_count > 0
                    documents_without_segments = 0 if doc_has_segments else 1
                else:
                    # 检查所有文档
                    segments_doc_ids = select(DocumentSegment.doc_id.distinct()).subquery()
                    documents_without_segments_stmt = select(DocumentMetadata).where(
                        ~DocumentMetadata.doc_id.in_(segments_doc_ids)
                    )
                    documents_without_segments_result = await session.execute(documents_without_segments_stmt)
                    documents_without_segments = len(documents_without_segments_result.scalars().all())
                
                # 生成报告
                integrity_report = {
                    "total_documents": metadata_count,
                    "total_segments": segment_count,
                    "orphaned_segments": orphaned_segments,
                    "documents_without_segments": documents_without_segments,
                    "data_consistent": orphaned_segments == 0 and documents_without_segments == 0
                }
                
                # 如果发现不一致，尝试修复
                if not integrity_report["data_consistent"]:
                    logger.warning(f"发现数据完整性问题: {integrity_report}")
                    
                    # 触发自动修复
                    fix_result = await self._fix_data_integrity(integrity_report)
                    integrity_report["fix_result"] = fix_result
                
                return integrity_report
                
        except Exception as e:
            logger.error(f"数据完整性验证失败: {str(e)}")
            return {
                "success": False,
                "error": str(e)
            }
    
    async def _fix_data_integrity(self, report: Dict[str, Any]) -> Dict[str, Any]:
        """
        修复数据完整性问题
        
        :param report: 完整性报告
        :return: 修复结果
        """
        try:
            async with self.async_session_factory() as session:
                async with session.begin():
                    fixed_orphaned = 0
                    fixed_empty_docs = 0
                    
                    # 处理孤立的分段
                    if report["orphaned_segments"] > 0:
                        # 使用子查询查找所有有效的doc_id
                        valid_doc_ids = select(DocumentMetadata.doc_id).subquery()
                        orphaned_segments_stmt = select(DocumentSegment).where(
                            ~DocumentSegment.doc_id.in_(valid_doc_ids)
                        )
                        orphaned_segments_result = await session.execute(orphaned_segments_stmt)
                        orphaned_segments = orphaned_segments_result.scalars().all()
                        
                        for segment in orphaned_segments:
                            await session.delete(segment)
                        
                        fixed_orphaned = len(orphaned_segments)
                        logger.info(f"删除了 {fixed_orphaned} 个孤立的文档分段")
                    
                    # 处理没有分段的文档
                    if report["documents_without_segments"] > 0:
                        # 这里可以根据需要实现不同的修复策略
                        # 例如：标记为删除，或者添加空分段等
                        logger.info(f"发现 {report['documents_without_segments']} 个没有分段的文档")
                    
                    return {
                        "success": True,
                        "fixed_orphaned_segments": fixed_orphaned,
                        "fixed_empty_documents": fixed_empty_docs
                    }
                    
        except Exception as e:
            logger.error(f"数据完整性修复失败: {str(e)}")
            return {
                "success": False,
                "error": str(e)
            }
    
    async def close(self):
        """
        关闭数据库连接，确保异步连接池被正确释放
        """
        try:
            if self.engine:
                # 确保异步连接池被正确关闭
                await self.engine.dispose()
                self.engine = None
                logger.info("数据库异步连接池已成功关闭和释放")
            if self.async_session_factory:
                # 清除会话工厂引用，避免内存泄漏
                self.async_session_factory = None
                logger.info("异步会话工厂已清除")
            return True
        except Exception as e:
            logger.error(f"关闭数据库连接失败: {str(e)}")
            return False


# 全局数据库管理器实例
database_manager = None


async def get_database_manager() -> DatabaseManager:
    """
    获取数据库管理器实例
    """
    global database_manager
    if database_manager is None:
        database_manager = DatabaseManager()
        # 异步创建数据库表
        await database_manager.create_tables()
    return database_manager